"""
python write_result_to_xlsx.py
"""
import os
import re
import json
import pandas as pd

# 定义目录路径
json_dir = '../fujian/fujian2/RandomForest'
output_file = '../fujian/result.xlsx'

# 获取所有的 json 文件
json_files = [f for f in os.listdir(json_dir) if f.endswith('.json')]

# 定义一个空字典存储所有类别的数据
data_dict = {}

# 设置目标日期范围
start_date = '2023/07/01'
end_date = '2023/09/30'
date_range = pd.date_range(start=start_date, end=end_date)
num_days = len(date_range)

# 使用正则表达式提取 category 编号并读取数据
for json_file in json_files:
    match = re.search(r'category(\d+)', json_file)
    if match:
        category_number = int(match.group(1))  # 提取 category 编号并转换为整数
        file_path = os.path.join(json_dir, json_file)
        
        with open(file_path, 'r') as f:
            data = json.load(f)
            # 仅提取指定日期范围的数据
            filtered_sales = [
                entry['predicted_sales']
                for entry in data
                if start_date <= entry['date'] <= end_date
            ]
            data_dict[category_number] = filtered_sales

# 创建 DataFrame
df = pd.DataFrame(index=[f'category{num}' for num in sorted(data_dict.keys())], columns=range(1, num_days + 1))

# 按 category 编号顺序填充数据
for category_number, sales_data in data_dict.items():
    if len(sales_data) != num_days:
        print(f"Warning: category{category_number} has an unexpected number of days.")
    df.loc[f'category{category_number}', :len(sales_data)] = sales_data

# 设置第一行和列的表头
df.columns = [f"{(start_date.split('/')[1] + '月')}{day}日" for day in range(1, num_days + 1)]
df.index.name = '类别'

# 导出到 Excel 文件
with pd.ExcelWriter(output_file) as writer:
    df.to_excel(writer, sheet_name='日销量预测结果', index=True)

print("数据已成功导出到 ../fujian/result.xlsx")
